package com.ps;


import jakarta.json.Json;
import jakarta.json.JsonObject;
import jakarta.websocket.*;
import jakarta.websocket.server.ServerEndpoint;
import java.io.StringReader;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;


@ServerEndpoint("/portscan")// 声明WebSocket服务端点路径
public class PortScanEndpoint {
    // 用于并发扫描的线程池
    private static final ExecutorService pool = Executors.newFixedThreadPool(20);

    /**
     * WebSocket连接建立时触发
     */
    @OnOpen
    public void onOpen(Session session) {
        System.out.println("Client connected: " + session.getId());
    }

    /**
     * 接收客户端消息时触发
     * 调用接口：jakarta.json.Json.createReader 解析JSON
     */
    @OnMessage
    public void onMessage(String msg, Session session) {
        // 解析 JSON 请求
        JsonObject req = Json.createReader(new StringReader(msg)).readObject();
        String host = req.getString("host");
        int from = req.getInt("fromPort");
        int to   = req.getInt("toPort");

        // 将扫描任务提交到线程池
        pool.submit(() -> scanPorts(host, from, to, session));
    }

    /**
     * WebSocket连接关闭时触发
     */
    @OnClose
    public void onClose(Session session) {
        System.out.println("Client disconnected: " + session.getId());
    }

    /**
     * 发生错误时触发
     */
    @OnError
    public void onError(Session session, Throwable thr) {
        thr.printStackTrace();
    }

    /**
     * 端口扫描核心方法
     */
    private void scanPorts(String host, int from, int to, Session session) {
        for (int port = from; port <= to; port++) {
            boolean open = isPortOpen(host, port, 200);
            // 构造返回 JSON
            JsonObject result = Json.createObjectBuilder()
                    .add("port", port)
                    .add("open", open)
                    .build();
            try {
                // 通过WebSocket实时发送结果
                session.getBasicRemote().sendText(result.toString());
            } catch (Exception e) {
                e.printStackTrace();
                break;
            }
        }
    }

    // 判断某端口是否开放
    private boolean isPortOpen(String host, int port, int timeoutMillis) {
        try (Socket socket = new Socket()) {
            socket.connect(new InetSocketAddress(host, port), timeoutMillis);
            return true;
        } catch (Exception ex) {
            return false;
        }
    }
}
